# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： img_io
Description :
Author : 'li'
date： 2022/1/19
-------------------------------------------------
Change Activity:
2022/1/19:
-------------------------------------------------
"""
import base64

import cv2
import numpy as np
from torch import Tensor

from ml import delete_file
from ...common import join_relative_path, create_dir
from ...common.misc.random_util import get_uuid_str


def read_img(path, to_rgb=False):
    """
    read img
    :param to_rgb:
    :param path:
    :return:
    """
    img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
    if img is None:
        return None
    if len(img.shape) == 2:
        return img
    img = img[:, :, :3]
    if to_rgb:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)
    return img


def write_img(img, path):
    """

    :param img:
    :param path:
    :return:
    """
    cv2.imencode('.jpg', img)[1].tofile(path)


def img_path_to_base64(img_path):
    """

    :param img_path:
    :return:
    """
    with open(img_path, "rb") as f:
        base64_data = base64.b64encode(f.read())
        base64_data = base64_data.decode('utf8')
        return base64_data


def image_to_base64(img):
    create_dir(join_relative_path('tmp_img'))
    path = join_relative_path('tmp_img/' + get_uuid_str() + '.jpg')
    write_img(img, path)
    base_str = img_path_to_base64(path)
    delete_file(path)
    return base_str


def base64_to_img(base64_encoding, des_img_path=None):
    """

    :param des_img_path:
    :param base64_encoding:
    :return:
    """
    base64_encoding = base64_encoding.encode('ascii')  # utf8
    base64_encoding = base64.b64decode(base64_encoding)
    img_np = np.frombuffer(base64_encoding, np.uint8)
    img = cv2.imdecode(img_np, cv2.COLOR_RGB2BGR)
    if des_img_path is not None:
        write_img(img, des_img_path)
    return img


def to_gray(copy_img):
    """
    image to gray
    :param copy_img:
    :return:
    """
    return cv2.cvtColor(copy_img, cv2.COLOR_BGR2GRAY)


def show_img(ori_img):
    if isinstance(ori_img, Tensor):
        img = ori_img.numpy()
    else:
        img = ori_img.copy()
    if np.max(img) <= 1:  # mask
        img = img * 255
    img = img.astype(np.uint8)
    # write_img(img, path='tmp.jpg')
    cv2.imshow('img', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def _main():
    bas = img_path_to_base64('1.jpg')
    print(bas)


if __name__ == '__main__':
    _main()
